import os
import cv2
import json
import numpy as np
from tqdm import tqdm

def coco_to_yolo_bbox(coco_bbox, img_w, img_h):
    x, y, w, h = coco_bbox
    x_center = (x + w / 2) / img_w
    y_center = (y + h / 2) / img_h
    w /= img_w
    h /= img_h
    return x_center, y_center, w, h

def get_cropped_bboxes_xywh(coco_bboxes, sub_coords, img_wh, slice_wh):
    iw, ih = img_wh
    sw, sh = slice_wh
    sx1, sy1, sx2, sy2 = sub_coords
    cropped_bboxes = []
    for cat_id, x, y, w, h in coco_bboxes:
        x1 = x * iw - (w * iw) / 2
        y1 = y * ih - (h * ih) / 2
        x2 = x * iw + (w * iw) / 2
        y2 = y * ih + (h * ih) / 2
        # Check intersection
        if x1 < sx2 and x2 > sx1 and y1 < sy2 and y2 > sy1:
            cx1 = max(x1, sx1) - sx1
            cy1 = max(y1, sy1) - sy1
            cx2 = min(x2, sx2) - sx1
            cy2 = min(y2, sy2) - sy1
            if cx2 > cx1 and cy2 > cy1:
                cx = (cx1 + cx2) / 2 / sw
                cy = (cy1 + cy2) / 2 / sh
                cw = (cx2 - cx1) / sw
                ch = (cy2 - cy1) / sh
                cropped_bboxes.append((cat_id, cx, cy, cw, ch))
    return cropped_bboxes

def split_coco_yolo(data_root, json_path, output_root, x, y, overlap):
    with open(json_path, 'r') as f:
        coco_data = json.load(f)

    img_id_to_annotations = {}
    for ann in coco_data['annotations']:
        img_id = ann['image_id']
        if img_id not in img_id_to_annotations:
            img_id_to_annotations[img_id] = []
        img_id_to_annotations[img_id].append(ann)

    id_to_filename = {img['id']: img['file_name'] for img in coco_data['images']}
    categories = {cat['id']: cat['name'] for cat in coco_data['categories']}

    os.makedirs(os.path.join(output_root, 'images'), exist_ok=True)
    os.makedirs(os.path.join(output_root, 'labels'), exist_ok=True)

    for img_id, file_name in tqdm(id_to_filename.items()):
        img_path = os.path.join(data_root, 'images', file_name)
        img = cv2.imread(img_path,cv2.IMREAD_UNCHANGED)
        if img is None:
            print(f"Warning: Failed to load {img_path}")
            continue

        ih, iw = img.shape[:2]
        bboxes = []
        for ann in img_id_to_annotations.get(img_id, []):
            coco_bbox = ann['bbox']
            x_c, y_c, w, h = coco_to_yolo_bbox(coco_bbox, iw, ih)
            bboxes.append((ann['category_id'], x_c, y_c, w, h))

        start_x = 0
        while start_x + x <= iw:
            start_y = 0
            while start_y + y <= ih:
                sx1, sy1 = start_x, start_y
                sx2, sy2 = start_x + x, start_y + y
                tile = img[sy1:sy2, sx1:sx2]
                tile_name = f"{os.path.splitext(file_name)[0]}_{sx1}_{sy1}_{x}_{y}_{overlap}.png"
                tile_path = os.path.join(output_root, 'images', tile_name)
                cv2.imwrite(tile_path, tile)

                cropped_labels = get_cropped_bboxes_xywh(bboxes, (sx1, sy1, sx2, sy2), (iw, ih), (x, y))
                label_path = os.path.join(output_root, 'labels', tile_name.replace('.png', '.txt'))
                with open(label_path, 'w') as f:
                    for label in cropped_labels:
                        f.write(" ".join([str(label[0])] + [f"{v:.6f}" for v in label[1:]]) + "\n")

                start_y += y - overlap
            start_x += x - overlap

if __name__ == "__main__":
    # data_root = "/home/lhx/data_zoo/truecolorseg"
    # json_path = "/home/lhx/data_zoo/truecolorseg/annotations.json"
    # output_root = "/home/lhx/data_zoo/splittruecolorseg"
    data_root = "/home/lhx/pcbsyn/truecolorwithdepthseg"
    json_path = "/home/lhx/data_zoo/truecolorseg/annotations.json"
    output_root = "/home/lhx/data_zoo/splittruecolorwithdepthseg"
    x = 1024
    y = 1024
    overlap = int(1024*0.1) 
    split_coco_yolo(data_root, json_path, output_root, x, y, overlap)
